/*
 * FEMCalc3D.cpp
 *
 *      Author: Jerome RD Soine and Christoph A Brand
 *      Institution: Schwarz goup, Institute for Theoretical Physics, Heidelberg University, Germany
 */

#include "../../../../include/base/substrate_models/FEM/FEMCalc.h"

FEMCalc::FEMCalc(double YOUNGS_MOD, double POISSON_RATIO, Point<3> P1,
		Point<3> P2, unsigned int N_REFS,
		vector<ForceTransmission*> FORCE_TRANS) {
	youngs_modulus = YOUNGS_MOD;
	poisson_ratio = POISSON_RATIO;
	force_trans = FORCE_TRANS;
	p1 = P1;
	p2 = P2;
	n_refs = N_REFS;
	generate_rectangular_mesh(p1, p2, force_trans, n_refs);

}

FEMCalc::~FEMCalc() {

}

SubstrateModel* FEMCalc::clone() const {
	SubstrateModel* sub_model = new FEMCalc(youngs_modulus, poisson_ratio, p1,
			p2, n_refs, force_trans);
	return sub_model;
}

void FEMCalc::run(vector<ForceTransmission*> FORCE_TRANS) {
	force_trans = FORCE_TRANS;
	project_traction(force_trans);
	setup_system();
	assemble_system();
	assemble_rhs(force_trans);
	apply_boundary_conditions();
	solve();
}

void FEMCalc::setup_system() {
	dof_handler->distribute_dofs(*fe);
	dof_handler->distribute_dofs(*fe_force);

	sparsity_pattern.reinit(dof_handler->n_dofs(), dof_handler->n_dofs(),
			dof_handler->max_couplings_between_dofs());

	DoFTools::make_sparsity_pattern(*dof_handler, sparsity_pattern);
	hanging_node_constraints.clear();
	DoFTools::make_hanging_node_constraints(*dof_handler,
			hanging_node_constraints);
	hanging_node_constraints.close();
	hanging_node_constraints.condense(sparsity_pattern);
	sparsity_pattern.compress();

	system_matrix.reinit(sparsity_pattern);

	solution.reinit(dof_handler->n_dofs());
	system_rhs.reinit(dof_handler->n_dofs());
}

void FEMCalc::assemble_system() {

	Timer timer;
	timer.start();
	/*
	 *Define Gauss quadrature for
	 *integration. Volume(2D,3D)
	 *and top surface.
	 */
	QGauss < 3 > quadrature_formula(2);
	/*
	 *FEM values Objects which provides
	 *the evaluation shape functions,
	 *gradients and jacobians
	 */
	FEValues < 3
			> fe_values(*fe, quadrature_formula,
					update_values | update_gradients | update_quadrature_points
							| update_JxW_values);
	/*
	 * Define constants needed for the
	 * numeric integration like DoF per cell,
	 * no of quadraature points (volume and face)
	 */
	const unsigned int dofs_per_cell = fe->dofs_per_cell;
	const unsigned int n_q_points = quadrature_formula.size();
	/*
	 * Define local matrix and rhs
	 */
	FullMatrix<double> cell_matrix(dofs_per_cell, dofs_per_cell);
	/*
	 * Define material constants in terms of
	 * lame coefficients (tranformation)
	 */
	double const_lambda = poisson_ratio * youngs_modulus / (1 + poisson_ratio)
			/ (1 - 2 * poisson_ratio);
	double const_mu = youngs_modulus / 2 / (1 + poisson_ratio);
	/*
	 * Additional stuff needed for the calculation.
	 */
	std::vector<unsigned int> local_dof_indices(dofs_per_cell);
	std::vector<double> lambda_values(n_q_points);
	std::vector<double> mu_values(n_q_points);
	ConstantFunction<3> lambda(const_lambda), mu(const_mu);
	//std::vector<Vector<double> > 	rhs_values (n_face_q_points, Vector<double>(spacedim));

	/*
	 * Define iteration object cell and last cell object
	 */
	DoFHandler<3>::active_cell_iterator cell = dof_handler->begin_active(),
			endc = dof_handler->end();

	/*
	 * Begin main integration loop
	 */
	for (; cell != endc; ++cell) {
		/*
		 * Reinit local matrix and rhs with null
		 */
		cell_matrix = 0;
		/*
		 * Associate fe_values object with current cell
		 */
		fe_values.reinit(cell);
		/*
		 * Get vector with material constants
		 */
		lambda.value_list(fe_values.get_quadrature_points(), lambda_values);
		mu.value_list(fe_values.get_quadrature_points(), mu_values);
		/*Calculate force vector with forces_function object*/

		//--------------------------------------
		//force_function.evaluate(fe_values.get_quadrature_points(),rhs_values);
		//--------------------------------------

		/*Loop over all DoFs of the current cell (local ennumeration)*/
		for (unsigned int i = 0; i < dofs_per_cell; ++i) {
			/*Get first component index*/
			const unsigned int component_i =
					fe->system_to_component_index(i).first;
			/*Loop over all DoFs of the current cell (local ennumeration)*/
			for (unsigned int j = 0; j < dofs_per_cell; ++j) {
				/*Get second component index*/
				const unsigned int component_j = fe->system_to_component_index(
						j).first;
				/*Loop over all qudrature points of the cell*/
				for (unsigned int q_point = 0; q_point < n_q_points;
						++q_point) {
					/*Matrix assembly (independent of the right hand side)*/
					cell_matrix(i, j) +=
							((fe_values.shape_grad(i, q_point)[component_i]
									* fe_values.shape_grad(j, q_point)[component_j]
									* lambda_values[q_point])
									+ (fe_values.shape_grad(i, q_point)[component_j]
											* fe_values.shape_grad(j, q_point)[component_i]
											* mu_values[q_point])
									+ ((component_i == component_j) ?
											(fe_values.shape_grad(i, q_point)
													* fe_values.shape_grad(j,
															q_point)
													* mu_values[q_point]) :
											0)) * fe_values.JxW(q_point);
				}
			}
		}

		/*Map the local system to the global one*/
		cell->get_dof_indices(local_dof_indices);
		for (unsigned int i = 0; i < dofs_per_cell; ++i) {
			for (unsigned int j = 0; j < dofs_per_cell; ++j)
				system_matrix.add(local_dof_indices[i], local_dof_indices[j],
						cell_matrix(i, j));
		}
	}

	/*Condense hanging_node_constraints*/
	hanging_node_constraints.condense(system_matrix);

	timer.stop();
}

void FEMCalc::assemble_rhs() {

	Timer timer;
	timer.start();

	system_rhs.reinit(dof_handler->n_dofs());
	solution.reinit(dof_handler->n_dofs());

	QGauss < 2 > face_quadrature_formula(2);

	FEFaceValues < 3
			> fe_face_values(*fe, face_quadrature_formula,
					update_values | update_q_points | update_JxW_values);

	FEFaceValues < 3
			> fe_force_values(*fe_force, face_quadrature_formula,
					update_values | update_q_points);

	const unsigned int n_face_q_points = face_quadrature_formula.size();
	const unsigned int dofs_per_face = fe->dofs_per_face;

	std::vector<unsigned int> local_dof_indices(dofs_per_face);
	std::vector<Vector<double> > rhs_values(n_face_q_points, Vector<double>(3));

	Vector<double> cell_rhs(dofs_per_face);

	std::vector<Tensor<1, 3> > sol_vec(n_face_q_points);

	DoFHandler<3>::active_cell_iterator cell = dof_handler->begin_active(),
			endc = dof_handler->end();

	const FEValuesExtractors::Vector displacements(0);

	for (; cell != endc; ++cell) {

		for (unsigned int face = 0; face < GeometryInfo < 3 > ::faces_per_cell;
				++face) {
			if (cell->face(face)->at_boundary()
					&& (cell->face(face)->boundary_indicator() == boundary_index)) {
				cell_rhs = 0;
				cell->face(face)->get_dof_indices(local_dof_indices);
				for (unsigned int i = 0; i < dofs_per_face; i++) {
					fe_face_values.reinit(cell, face);
					fe_force_values.reinit(cell, face);
					fe_force_values[displacements].get_function_values(
							force_solution, sol_vec);
					unsigned int j = fe->face_to_cell_index(i, face);

					const unsigned int component_i =
							fe->system_to_component_index(j).first;

					for (unsigned int q_point = 0; q_point < n_face_q_points;
							++q_point)
						cell_rhs(i) += fe_face_values.shape_value(j, q_point)
								* sol_vec[q_point][component_i]
								* fe_face_values.JxW(q_point);

					system_rhs(local_dof_indices[i]) += cell_rhs(i);
				}
			}
		}

	}

	hanging_node_constraints.condense(system_rhs);
	timer.stop();

}

void FEMCalc::assemble_rhs(vector<ForceTransmission*> force_transmission) {

	Timer timer;
	timer.start();

	system_rhs.reinit(dof_handler->n_dofs());
	solution.reinit(dof_handler->n_dofs());

	QGauss < 2 > face_quadrature_formula(2);

	FEFaceValues < 3
			> fe_face_values(*fe, face_quadrature_formula,
					update_values | update_q_points | update_JxW_values);

	FEFaceValues < 3
			> fe_force_values(*fe_force, face_quadrature_formula,
					update_values | update_q_points);

	const unsigned int n_face_q_points = face_quadrature_formula.size();
	const unsigned int dofs_per_face = fe->dofs_per_face;

	std::vector<unsigned int> local_dof_indices(dofs_per_face);
	std::vector<Vector<double> > rhs_values(n_face_q_points, Vector<double>(3));

	Vector<double> cell_rhs(dofs_per_face);

	std::vector<Tensor<1, 3> > sol_vec(n_face_q_points);

	DoFHandler<3>::active_cell_iterator cell = dof_handler->begin_active(),
			endc = dof_handler->end();

	const FEValuesExtractors::Vector displacements(0);

	for (; cell != endc; ++cell) {
		for (unsigned int face = 0; face < GeometryInfo < 3 > ::faces_per_cell;
				++face) {
			if (cell->face(face)->at_boundary()
					&& (cell->face(face)->boundary_indicator() == boundary_index)) {
				cell_rhs = 0;
				cell->face(face)->get_dof_indices(local_dof_indices);
				for (unsigned int i = 0; i < dofs_per_face; i++) {
					fe_face_values.reinit(cell, face);
					fe_force_values.reinit(cell, face);
					fe_force_values[displacements].get_function_values(
							force_solution, sol_vec);
					unsigned int j = fe->face_to_cell_index(i, face);

					const unsigned int component_i =
							fe->system_to_component_index(j).first;

					for (unsigned int q_point = 0; q_point < n_face_q_points;
							++q_point)
						cell_rhs(i) += fe_face_values.shape_value(j, q_point)
								* sol_vec[q_point][component_i]
								* fe_face_values.JxW(q_point);

					system_rhs(local_dof_indices[i]) += cell_rhs(i);
				}
			}
		}
	}

	hanging_node_constraints.condense(system_rhs);

	timer.stop();
}

void FEMCalc::apply_boundary_conditions() {
	Timer timer;
	timer.start();

	//std::cout << "FEMCalc apply Boundary Conditions ... \n";

	/* Interpolate boundary conditions the boundaries.
	 * Here we have to destinguish the 2D and 3D case.
	 * In 2D we applie zero boundray conditions to all
	 * boundaries. In 3D we apply only to the non top
	 * boundraies zero boundray conditions.
	 */
	std::map<unsigned int, double> boundary_values;

	// Set more BCs

	VectorTools::interpolate_boundary_values(*dof_handler, boundary_index - 1,
			ZeroFunction < 3 > (3), boundary_values);

	MatrixTools::apply_boundary_values(boundary_values, system_matrix, solution,
			system_rhs);

	timer.stop();

}

void FEMCalc::solve() {

	Timer timer;
	timer.start();

	SolverControl solver_control(10000, 1e-6);
	SolverCG<> cg(solver_control);
	PreconditionSSOR<> preconditioner;
	preconditioner.initialize(system_matrix, 1.2);
	cg.solve(system_matrix, solution, system_rhs, preconditioner);
	hanging_node_constraints.distribute(solution);
	timer.stop();
}
